package com.lingjtx.common.thread.core;

import com.alibaba.nacos.shaded.com.google.common.util.concurrent.ThreadFactoryBuilder;
import com.lingjtx.common.core.util.StringUtil;
import com.lingjtx.common.thread.config.ExecutorConfig;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.concurrent.*;

public class ThreadPoolBuilder {

    private int corePoolSize = Runtime.getRuntime().availableProcessors();
    private int maxPoolSize = corePoolSize * 2;
    private int queueCapacity = 100;
    private int keepAliveSeconds = 60;
    private String rejectedPolicy = "CallerRunsPolicy";

    public ThreadPoolBuilder corePoolSize(int corePoolSize) {
        this.corePoolSize = corePoolSize;
        return this;
    }

    public ThreadPoolBuilder maxPoolSize(int maxPoolSize) {
        this.maxPoolSize = maxPoolSize;
        return this;
    }

    public ThreadPoolBuilder queueCapacity(int queueCapacity) {
        this.queueCapacity = queueCapacity;
        return this;
    }

    public ThreadPoolBuilder keepAliveSeconds(int keepAliveSeconds) {
        this.keepAliveSeconds = keepAliveSeconds;
        return this;
    }

    // 设置拒绝策略
    public ThreadPoolBuilder rejectedPolicy(String rejectedPolicy) {
        this.rejectedPolicy = rejectedPolicy;
        return this;
    }

    // 创建 ThreadPoolExecutor
    public ThreadPoolExecutor build(String name) {
        return new ThreadPoolExecutor(corePoolSize, maxPoolSize, keepAliveSeconds, TimeUnit.SECONDS,
                new LinkedBlockingQueue<>(queueCapacity),
                r -> new Thread(r, name + "-thread-" + StringUtil.getUuid()),
                getRejectedExecutionHandler(rejectedPolicy));
    }

    // 创建 Spring 的 ThreadPoolTaskExecutor
    public static ThreadPoolTaskExecutor buildSpringExecutor(String name, ThreadPoolExecutor executor, ExecutorConfig.PoolConfig config) {
        ThreadPoolTaskExecutor springExecutor = new ThreadPoolTaskExecutor() {
            @Override
            protected ExecutorService initializeExecutor(ThreadFactory threadFactory, RejectedExecutionHandler handler) {
                return executor; // 强制使用我们构建的 executor (共享底层线程池)
            }
        };
        springExecutor.setCorePoolSize(config.getCorePoolSize());
        springExecutor.setMaxPoolSize(config.getMaxPoolSize());
        springExecutor.setQueueCapacity(config.getQueueCapacity());
        springExecutor.setKeepAliveSeconds(config.getKeepAliveSeconds());
        springExecutor.setThreadNamePrefix(name + "-thread-");
        springExecutor.setRejectedExecutionHandler(getRejectedExecutionHandler(config.getRejectedPolicy()));
        return springExecutor;
    }

    public static ThreadPoolExecutor build(String name, ExecutorConfig.PoolConfig config) {
        BlockingQueue<Runnable> queue = new LinkedBlockingQueue<>(config.getQueueCapacity());
        ThreadFactory factory = new ThreadFactoryBuilder()
                .setNameFormat(name + "-thread-%d")
                .build();
        return new ThreadPoolExecutor(
                config.getCorePoolSize(),
                config.getMaxPoolSize(),
                config.getKeepAliveSeconds(),
                TimeUnit.SECONDS,
                queue,
                factory,
                getRejectedExecutionHandler(config.getRejectedPolicy())
        );
    }

    public static RejectedExecutionHandler getRejectedExecutionHandler(String policy) {
        return switch (policy) {
            case "AbortPolicy" -> new ThreadPoolExecutor.AbortPolicy();
            case "DiscardPolicy" -> new ThreadPoolExecutor.DiscardPolicy();
            case "DiscardOldestPolicy" -> new ThreadPoolExecutor.DiscardOldestPolicy();
            default -> new ThreadPoolExecutor.CallerRunsPolicy();
        };
    }
}
